import PIL.Image as Image
import io
import os
import numpy as np

import tensorflow as tf

from object_tutorial import tfrecord_decoder

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

filename = 'D:/Work/workspace/helmet_fire.record'


def convert_tensor_dataset_to_box_and_text(e_ymin, e_xmin, e_ymax, e_xmax, e_text):
    length = e_ymax[1].__len__()
    boxes = []
    texts = []
    for i in range(length):
        boxes.append((e_ymin[1][i], e_xmin[1][i], e_ymax[1][i], e_xmax[1][i]))
        texts.append(e_text[1][i])
    return boxes, texts


dataset = tf.data.TFRecordDataset(filename)
dataset = dataset.shuffle(buffer_size=1000)
dataset = dataset.repeat(100)
iterator = dataset.make_one_shot_iterator()
element = iterator.get_next()
with tf.Session() as sess:
    for i in range(10):
        example = sess.run(element)
        image, height, width, filename, xmin, xmax, ymin, ymax, text \
            = tfrecord_decoder.decode_example(example)
        e_image, e_height, e_width, e_filename, e_xmin, e_xmax, e_ymin, e_ymax, e_text = \
            sess.run([image, height, width, filename, xmin, xmax, ymin, ymax, text])
        box = e_ymin, e_xmin, e_ymax, e_xmax
        image_io = io.BytesIO(e_image)
        image = Image.open(image_io)
        print(image)
        np.array(image.getdata()).reshape((720, 1280, 3)).astype(np.uint8)
